package com.project.util;


import java.rmi.UnexpectedException;
import java.util.concurrent.atomic.AtomicLong;

public class SnowflakeDistributedId {

    /**
     * Snowflake算法中第三部分长度，即数据中心和工作机器ID总共占位长度
     */
    private static final long DATA_CENTER_AND_WORKER_ID_BITS = 10;

    /**
     * Snowflake算法中第四部分长度，即自增序列占位长度
     */
    private static final long AUTO_INCREMENT_SEQUENCE_BITS = 12;

    /**
     * 自增序列最大值
     */
    private static final long MAX_SEQUENCE = 4095;

    /**
     * 开始时间戳
     */
    private final long epoch;

    /**
     * 数据中心ID
     */
    private final long dataCenterId;

    /**
     * 机器ID占位长度
     */
    private final long workerIdBits;

    /**
     * 机器ID
     */
    private final long workerId;

    /**
     * 保存上一次生成ID的时间戳
     */
    private long lastTimestamp;

    /**
     * 分布式ID自增序列
     */
    private AtomicLong autoIncrementSequence;

    /**
     * @param dataCenterIdBits 数据中心ID占位长度
     * @param dataCenterId     数据中心ID
     * @param workerId         工作机器ID
     */
    public SnowflakeDistributedId(long epoch, long dataCenterIdBits, long dataCenterId, long workerId) {
        this.epoch = epoch;
        this.dataCenterId = validateDataCenterId(dataCenterIdBits, dataCenterId);
        workerIdBits = DATA_CENTER_AND_WORKER_ID_BITS - dataCenterIdBits;
        this.workerId = validateWorkerId(workerId);
        this.lastTimestamp = -1L;
        this.autoIncrementSequence = new AtomicLong(0);
    }

    /**
     * 初始化数据中心ID
     *
     * @param dataCenterIdBits 数据中心ID占位长度
     * @param dataCenterId     数据中心ID
     * @return 校验通过的数据中心ID
     */
    private long validateDataCenterId(long dataCenterIdBits, long dataCenterId) {
        if (dataCenterIdBits < 0 || dataCenterIdBits >= DATA_CENTER_AND_WORKER_ID_BITS) {
            throw new IllegalArgumentException("Data center ID bits must be in [0, 10)!");
        }
        if (dataCenterIdBits > 0) {
            // 支持的最大数据中心 ID
            long maxDataCenterId = ~(-1 << dataCenterIdBits);
            if (dataCenterId < 0 || dataCenterId > maxDataCenterId) {
                throw new IllegalArgumentException("Data center ID must be in [0, " + maxDataCenterId + "]!");
            }
            return dataCenterId;
        }
        return -1;
    }

    /**
     * 初始化工作机器ID
     *
     * @param workerId 工作机器ID
     * @return 校验通过的工作机器ID
     */
    private long validateWorkerId(long workerId) {
        // 支持的最大机器ID
        long maxWorkerId = ~(-1 << this.workerIdBits);
        if (workerId < 0 || workerId > maxWorkerId) {
            throw new IllegalArgumentException("Worker ID must be in [0, " + maxWorkerId + "]!");
        }
        return workerId;
    }

    /**
     * 生成分布式ID
     *
     * @return long类型ID
     * @throws UnexpectedException 如果系统时间回退则抛出此异常
     */
    public synchronized long generate() throws UnexpectedException {
        long currentTimestamp = System.currentTimeMillis();
        // 如果当前时间小于上一次ID生成时间，说明系统时间回退
        if (currentTimestamp < lastTimestamp) {
            throw new UnexpectedException("System clock moved backward, refused to generate ID!");
        }
        long currentSequence;
        if (currentTimestamp == lastTimestamp) {
            // 如果当前时间等于上一次ID生成时间，获取自增序列值后加1
            currentSequence = autoIncrementSequence.getAndIncrement();
            // 如果获取的自增序列值大于允许的最大值
            if (currentSequence > MAX_SEQUENCE) {
                // 等待到下一毫秒
                currentTimestamp = block(currentTimestamp);
                // 更新时间戳
                lastTimestamp = currentTimestamp;
                // 重新获取自增序列值
                currentSequence = resetAutoIncrementSequence();
            }
        } else {
            // 如果当前时间大于上一次ID生成时间，重置自增序列并获取自增序列值后加1
            currentSequence = resetAutoIncrementSequence();
            // 更新时间戳
            lastTimestamp = currentTimestamp;
        }
        // 时间戳左移
        long id = (currentTimestamp - epoch) << (DATA_CENTER_AND_WORKER_ID_BITS + AUTO_INCREMENT_SEQUENCE_BITS);
        if (dataCenterId != -1) {
            // 数据中心ID左移
            id = id | (this.dataCenterId << (workerIdBits + AUTO_INCREMENT_SEQUENCE_BITS));
        }
        return id | (this.workerId << AUTO_INCREMENT_SEQUENCE_BITS) | currentSequence;
    }

    /**
     * 重置自增序列
     *
     * @return 自增序列值
     */
    private synchronized long resetAutoIncrementSequence() {
        autoIncrementSequence = new AtomicLong(0);
        return autoIncrementSequence.getAndIncrement();
    }

    /**
     * 阻塞至下一毫秒
     *
     * @param timestamp 当前时间戳
     * @return 下一毫秒时间戳
     */
    private long block(long timestamp) {
        long currentTimestamp = System.currentTimeMillis();
        while (currentTimestamp <= timestamp) {
            currentTimestamp = System.currentTimeMillis();
        }
        return currentTimestamp;
    }
}
